
// llvmir.cpp
#include "llvmir.h"
#include <memory>
#include <cctype>

namespace LLVMIR
{
    RESOURCE::Uint_Allocator label_allocator(0);

    std::string llvm_program;
    int main_end_index = 0;
    int get_type_size(const std::string &type)
    {
        if (type == "i32")
        {
            return 4;
        }
        else if (type == "i64")
        {
            return 8;
        }
        else if (type == "float")
        {
            return 4;
        }
        else if (type == "double")
        {
            return 8;
        }
        else
        {
            return 0;
        }
    }

    void add_llvm_code(const std::string &code)
    {
        llvm_program = llvm_program + code;
    }

    void print_llvm_code()
    {
        std::cout << llvm_program << std::endl;
    }

    void clear_llvm_code()
    {
        llvm_program.clear();
    }

    void write_llvm_code_to_file(const std::string &file_name)
    {
        std::ofstream out(file_name);
        out << llvm_program;
        out.close();
    }

    void const_decl(std::string name, int value, int level, bool is_globle)
    {
        std::string code = " @" + name + " = constant i32 " + std::to_string(value) + "\n";
        add_llvm_code(code);
    }

    void alloc_stack_space(std::string name, std::string type)
    {
        int size = get_type_size(type);
        std::string code = " %" + name + " = alloca " + type + ", align " + std::to_string(size) + "\n";
        add_llvm_code(code);
    }

    void store(std::string name, std::string type, int value)
    {
        std::string code = " store " + type + " " + std::to_string(value) + ", " + type + "* %" + name + ", align " + std::to_string(get_type_size(type)) + "\n";
        add_llvm_code(code);
    }

    void load(std::string name, std::string type, std::string result)
    {
        std::string code = " " + result + " = load " + type + ", " + type + "* %" + name + ", align " + std::to_string(get_type_size(type)) + "\n";
        add_llvm_code(code);
    }

    void add(std::string type, std::string result, std::string op1, std::string op2)
    {
        std::string code = " " + result + " = add " + type + " " + op1 + ", " + op2 + "\n";
        add_llvm_code(code);
    }

    void sub(std::string type, std::string result, std::string op1, std::string op2)
    {
        std::string code = " " + result + " = sub " + type + " " + op1 + ", " + op2 + "\n";
        add_llvm_code(code);
    }

    void mul(std::string type, std::string result, std::string op1, std::string op2)
    {
        std::string code = " " + result + " = mul " + type + " " + op1 + ", " + op2 + "\n";
        add_llvm_code(code);
    }

    void div(std::string type, std::string result, std::string op1, std::string op2)
    {
        std::string code = " " + result + " = sdiv " + type + " " + op1 + ", " + op2 + "\n";
        add_llvm_code(code);
    }

    void icmp(std::string type, std::string result, std::string op1, std::string op2)
    {
        std::string code = " " + result + " = icmp eq " + type + " " + op1 + ", " + op2 + "\n";
        add_llvm_code(code);
    }

    void br(std::string label)
    {
        std::string code = " br label %" + label + "\n";
        add_llvm_code(code);
    }

    void define_function(std::string procedure_name)
    {
        std::string name = procedure_name.substr(10);
        std::string code = " define void @" + name + "()\n";
        add_llvm_code(code);
    }

    void define_main_function()
    {
        std::string code = "define void @main()\n";
        add_llvm_code(code);
    }

    void const_array_decl(std::string name, std::vector<int> values, int size, int level, bool is_globle)
    {
        std::string code;
        if (is_globle)
        {
            code = " @" + name + " = constant [";
        }
        else
        {
            code = " %" + name + " = alloca [";
        }
        code += std::to_string(size);
        code += " x i32] [";
        for (size_t i = 0; i < values.size(); ++i)
        {
            code += std::to_string(values[i]);
            if (i != values.size() - 1)
            {
                code += ", ";
            }
        }
        for (size_t i = values.size(); i < size; ++i)
        {
            code += ", 0";
        }
        code += "]\n";
        add_llvm_code(code);
    }
    void var_decl(std::string name, int value, int level, bool is_globle)
    {
        if (level != 0 && is_globle)
        {
            std::cout << "error: globle var in block" << std::endl;
            return;
        }
        std::string code;
        if (is_globle)
        {
            code = " @" + name + " = global i32 " + std::to_string(value) + "\n";
        }
        else
        {
            code = " %" + name + " = alloca i32\n";
            code += " store i32 " + std::to_string(value) + ", i32* %" + name + "\n";
        }
        add_llvm_code(code);
    }

    void var_decl_array(std::string name, std::vector<int> values, int size, int level, bool is_globle)
    {
        if (level != 0 && is_globle)
        {
            std::cout << "error: global var in block" << std::endl;
            return;
        }

        std::string code;

        if (is_globle)
        {
            code = " @" + name + " = global [" + std::to_string(size) + " x i32] [";

            for (size_t i = 0; i < values.size(); ++i)
            {
                code += std::to_string(values[i]);
                if (i != values.size() - 1)
                {
                    code += ", ";
                }
            }

            for (size_t i = values.size(); i < size; ++i)
            {
                code += ", 0";
            }

            code += "]\n";
        }
        else
        {
            code = " %" + name + " = alloca [" + std::to_string(size) + " x i32]\n";
            for (size_t i = 0; i < values.size(); ++i)
            {
                code += " store i32 " + std::to_string(values[i]) + ", i32* getelementptr inbounds ([" + std::to_string(size) + " x i32], [" + std::to_string(size) + " x i32]* %" + name + ", i32 0, i32 " + std::to_string(i) + ")\n";
            }

            for (size_t i = values.size(); i < size; ++i)
            {
                code += " store i32 0, i32* getelementptr inbounds ([" + std::to_string(size) + " x i32], [" + std::to_string(size) + " x i32]* %" + name + ", i32 0, i32 " + std::to_string(i) + ")\n";
            }
        }
        add_llvm_code(code);
    }

    void write(std::string name, bool is_globle)
    {
        int num = label_allocator.allocate();
        std::string call_label = "%call" + std::to_string(num);

        std::string load_code;
        if (is_globle)
        {
            load_code = "%" + name + "_val = load i32, i32* @" + name + "\n";
        }
        else
        {
            load_code = "%" + name + "_val = load i32, i32* %" + name + "\n";
        }
        add_llvm_code(load_code);

        std::string code = call_label + " = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([4 x i8], [4 x i8]* @.str, i32 0, i32 0), i32 %" + name + "_val)\n";

        add_llvm_code(code);
    }

    void read(std::string name)
    {
        int num = label_allocator.allocate();
        std::string call_label = "%call" + std::to_string(num);

        std::string code = call_label + " = call i32 (i8*, ...) @scanf(i8* getelementptr inbounds ([3 x i8], [3 x i8]* @.str.scanf, i32 0, i32 0), i32* %" + name + ")\n";

        add_llvm_code(code);
    }

    void add_left_brace(bool front)
    {
        if (front)
        {
            add_llvm_code("{\n");
        }
        else
        {
            add_llvm_code("{\n");
        }
    }

    void add_right_brace(bool front)
    {
        if (front)
        {
            add_llvm_code("}\n");
        }
        else
        {
            add_llvm_code("}\n");
        }
    }

    void call(std::string name)
    {
        std::string code = " call void @" + name + "()\n";
        add_llvm_code(code);
    }

    void insert_in_main(std::string code)
    {
        llvm_program.insert(main_end_index, code);
        main_end_index += code.size();
    }

    void add_left_brace()
    {
        add_llvm_code("{\n");
    }

    void add_right_brace()
    {
        add_llvm_code("}\n");
    }

    void const_decl_array(std::string name, std::vector<int> values, int size, bool is_globle)
    {
        std::string code;

        if (is_globle)
        {
            code = " @" + name + " = constant [" + std::to_string(size) + " x i32] [";
            for (size_t i = 0; i < values.size(); ++i)
            {
                code += "i32 " + std::to_string(values[i]);
                if (i != values.size() - 1)
                {
                    code += ", ";
                }
            }
            for (size_t i = values.size(); i < size; ++i)
            {
                code += ", i32 0";
            }
            code += "]\n";
        }
        else
        {
            std::cout << "error: local constant arrays are not supported in LLVM IR" << std::endl;
            return;
        }

        add_llvm_code(code);
    }

    void add_printf_and_scanf()
    {
        std::string code = " @.str = private unnamed_addr constant [4 x i8] c\"%d\\0A\\00\"\n";
        code += " @.str.scanf = private unnamed_addr constant [3 x i8] c\"%d\\00\"\n";
        code += "declare i32 @scanf(i8*, ...)\n";
        code += "declare i32 @printf(i8*, ...)\n";
        add_llvm_code(code);
    }

    void add_return_void()
    {
        add_llvm_code(" ret void\n");
    }

    void add_icmp(std::string condition, std::string lhs, std::string rhs,std::string name)
    {
        int num = label_allocator.allocate();
        bool is_lhs_ident = isdigit(lhs[0]);
        bool is_rhs_ident = isdigit(rhs[0]);
        std::string lhs_tmp;
        std::string rhs_tmp;
        if(!is_lhs_ident)
        {
            lhs_tmp = "%" + lhs + std::to_string(num);
            add_llvm_code(lhs_tmp + "= load i32 ,i32* %"+lhs+"\n");
            lhs ="%" + lhs + std::to_string(num);
        }
        
        if(!is_rhs_ident)
        {
            rhs_tmp = "%" +rhs + std::to_string(num);
            add_llvm_code(rhs_tmp + "= load i32 ,i32* %"+rhs+"\n");
            rhs = "%" + rhs + std::to_string(num) ;
        }
        name = name + "_";
        add_llvm_code(" %cmp = icmp sgt i32 " + lhs + ", " + rhs + "\n");
        add_llvm_code(" br i1 %cmp, label %" + name + "true, label %" + name + "false\n");
    }

    void add_if_true(std::string name)
    {
        add_llvm_code(name+":"+"\n");
    }

    void add_if_false(std::string name)
    {
        add_llvm_code(name+":"+"\n");
    }
}
